package io.keyss.view_record.recording

import android.graphics.Bitmap
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaCodecList
import android.media.MediaFormat
import android.media.MediaMuxer
import android.util.Log
import java.io.File

class BitmapToMp4Encoder(
    private val width: Int,
    private val height: Int,
    private val frameRate: Int = 30,
    private val bitRate: Int = 2_000_000 // 默认 2Mbps
) {
    private var mediaCodec: MediaCodec? = null
    private var mediaMuxer: MediaMuxer? = null
    private var trackIndex = -1
    private var isMuxerStarted = false
    private var presentationTimeUs: Long = 0
    private var outputFile: File? = null
    private var startTime: Long = -1
    private var isReleasing = false

    // 初始化编码器
    fun prepare(outputPath: String) {
        outputFile = File(outputPath)
        if (outputFile!!.exists()) {
            outputFile!!.delete()
        }
        outputFile!!.createNewFile()

        try {
            mediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC).apply {
                val format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height).apply {
                    setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar)
                    setInteger(MediaFormat.KEY_BIT_RATE, bitRate)
                    setInteger(MediaFormat.KEY_FRAME_RATE, frameRate)
                    setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1)
                    setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, width * height * 2)
                }
                configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
                start()
            }

            mediaMuxer = MediaMuxer(outputFile!!.path, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)

        } catch (e: Exception) {
            Log.e("Encoder", "Failed to initialize encoder: ${e.message}", e)
            throw e
        }
    }

    // 编码单帧
    fun encodeFrame(bitmap: Bitmap) {
        if (isReleasing) return

        if (startTime == -1L) {
            startTime = System.nanoTime() / 1000
        }

        val currentTimeUs = System.nanoTime() / 1000
        presentationTimeUs = currentTimeUs - startTime

        try {
            // 确保bitmap尺寸与编码器尺寸匹配
            val scaledBitmap = if (bitmap.width != width || bitmap.height != height) {
                Bitmap.createScaledBitmap(bitmap, width, height, true)
            } else {
                bitmap
            }

            val yuvData = convertBitmapToYUV420(scaledBitmap)

            // 如果bitmap被缩放，需要回收
            if (scaledBitmap != bitmap) {
                scaledBitmap.recycle()
            }

            mediaCodec?.let { codec ->
                val inputBufferId = codec.dequeueInputBuffer(10_000)
                if (inputBufferId >= 0) {
                    val inputBuffer = codec.getInputBuffer(inputBufferId)
                    if (inputBuffer != null) {
                        if (yuvData.size > inputBuffer.capacity()) {
                            Log.e("Encoder", "Buffer overflow: yuvData size ${yuvData.size} > buffer capacity ${inputBuffer.capacity()}")
                            return
                        }

                        inputBuffer.clear()
                        inputBuffer.put(yuvData)
                        codec.queueInputBuffer(
                            inputBufferId,
                            0,
                            yuvData.size,
                            presentationTimeUs,
                            0
                        )
                    }
                }

                val bufferInfo = MediaCodec.BufferInfo()
                var outputBufferId = codec.dequeueOutputBuffer(bufferInfo, 10_000)

                while (outputBufferId >= 0) {
                    val outputBuffer = codec.getOutputBuffer(outputBufferId)
                    if (outputBuffer != null) {
                        if (!isMuxerStarted && bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG != 0) {
                            codec.releaseOutputBuffer(outputBufferId, false)
                        } else {
                            if (!isMuxerStarted) {
                                trackIndex = mediaMuxer?.addTrack(codec.outputFormat) ?: -1
                                mediaMuxer?.start()
                                isMuxerStarted = true
                            }

                            if (bufferInfo.size > 0) {
                                outputBuffer.position(bufferInfo.offset)
                                outputBuffer.limit(bufferInfo.offset + bufferInfo.size)
                                mediaMuxer?.writeSampleData(trackIndex, outputBuffer, bufferInfo)
                            }
                            codec.releaseOutputBuffer(outputBufferId, false)
                        }
                    }
                    outputBufferId = codec.dequeueOutputBuffer(bufferInfo, 0)
                }
            }
        } catch (e: Exception) {
            Log.e("Encoder", "Encode frame error: ${e.message}", e)
        }
    }

    // 结束编码
    fun release() {
        if (isReleasing) return
        isReleasing = true

        try {
            mediaCodec?.let { codec ->
                // 标记结束
                val inputBufferId = codec.dequeueInputBuffer(10_000)
                if (inputBufferId >= 0) {
                    codec.queueInputBuffer(
                        inputBufferId,
                        0,
                        0,
                        presentationTimeUs,
                        MediaCodec.BUFFER_FLAG_END_OF_STREAM
                    )
                }

                // 处理剩余数据
                val bufferInfo = MediaCodec.BufferInfo()
                var outputBufferId = codec.dequeueOutputBuffer(bufferInfo, 10_000)
                while (outputBufferId >= 0) {
                    val outputBuffer = codec.getOutputBuffer(outputBufferId)
                    if (outputBuffer != null && isMuxerStarted) {
                        if (bufferInfo.size > 0) {
                            outputBuffer.position(bufferInfo.offset)
                            outputBuffer.limit(bufferInfo.offset + bufferInfo.size)
                            mediaMuxer?.writeSampleData(trackIndex, outputBuffer, bufferInfo)
                        }
                    }
                    codec.releaseOutputBuffer(outputBufferId, false)
                    if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM != 0) {
                        break
                    }
                    outputBufferId = codec.dequeueOutputBuffer(bufferInfo, 10_000)
                }

                codec.stop()
                codec.release()
            }
            mediaCodec = null

            mediaMuxer?.apply {
                if (isMuxerStarted) {
                    stop()
                    release()
                }
            }
            mediaMuxer = null
            isMuxerStarted = false

        } catch (e: Exception) {
            Log.e("Encoder", "Release error: ${e.message}", e)
        } finally {
            isReleasing = false
        }
    }

    private  val MIME_TYPE = "video/avc" // H.264 Advanced Video Coding

    // Bitmap转YUV420（关键方法）
    private fun convertBitmapToYUV420(bitmap: Bitmap): ByteArray {
        val width = bitmap.width
        val height = bitmap.height
        val ySize = width * height
        val uvSize = width * height / 4
        val yuv = ByteArray(ySize + uvSize * 2)

        val argb = IntArray(width * height)
        bitmap.getPixels(argb, 0, width, 0, 0, width, height)

        // 使用优化的YUV转换算法
        var yIndex = 0
        var uvIndex = ySize

        // 先转换Y分量
        for (j in 0 until height) {
            for (i in 0 until width) {
                val pixel = argb[j * width + i]
                val r = (pixel shr 16) and 0xff
                val g = (pixel shr 8) and 0xff
                val b = pixel and 0xff

                // Y = 0.299R + 0.587G + 0.114B
                yuv[yIndex++] = ((0.299f * r + 0.587f * g + 0.114f * b).toInt() and 0xff).toByte()
            }
        }

        // 转换UV分量
        for (j in 0 until height step 2) {
            for (i in 0 until width step 2) {
                var r = 0
                var g = 0
                var b = 0

                // 计算2x2块的RGB平均值
                for (y in 0..1) {
                    for (x in 0..1) {
                        if (j + y < height && i + x < width) {
                            val pixel = argb[(j + y) * width + (i + x)]
                            r += (pixel shr 16) and 0xff
                            g += (pixel shr 8) and 0xff
                            b += pixel and 0xff
                        }
                    }
                }

                // 计算平均值
                r /= 4
                g /= 4
                b /= 4

                // U = -0.147R - 0.289G + 0.436B + 128
                // V = 0.615R - 0.515G - 0.100B + 128
                val u = ((-0.147f * r - 0.289f * g + 0.436f * b + 128).toInt() and 0xff).toByte()
                val v = ((0.615f * r - 0.515f * g - 0.100f * b + 128).toInt() and 0xff).toByte()

                yuv[uvIndex++] = u
                yuv[uvIndex++] = v
            }
        }

        return yuv
    }

    private fun selectCodec(mimeType: String): MediaCodecInfo? {
        val numCodecs = MediaCodecList.getCodecCount()
        for (i in 0 until numCodecs) {
            val codecInfo = MediaCodecList.getCodecInfoAt(i)
            if (!codecInfo.isEncoder) {
                continue
            }
            val types = codecInfo.supportedTypes
            for (j in types.indices) {
                if (types[j].equals(mimeType, ignoreCase = true)) {
                    return codecInfo
                }
            }
        }
        return null
    }

    private fun selectColorFormat(
        codecInfo: MediaCodecInfo,
        mimeType: String
    ): Int {
        val capabilities = codecInfo
            .getCapabilitiesForType(mimeType)
        for (i in capabilities.colorFormats.indices) {
            val colorFormat = capabilities.colorFormats[i]
            if (isRecognizedFormat(colorFormat)) {
                return colorFormat
            }
        }
        return 0 // not reached
    }
    private fun isRecognizedFormat(colorFormat: Int): Boolean {
        return when (colorFormat) {
            MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420PackedPlanar, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420PackedSemiPlanar, MediaCodecInfo.CodecCapabilities.COLOR_TI_FormatYUV420PackedSemiPlanar -> true
            else -> false
        }
    }
}
